from crewai import Agent, Crew, Process, Task
from crewai.project import CrewBase, agent, crew, task
from inventory_management.tools.database_tool import DatabaseTool
from inventory_management.tools.nlp_tool import NLPTool
from inventory_management.tools.inventory_tool import InventoryTool
from langchain_openai import ChatOpenAI
from langchain.tools import Tool

# Uncomment the following line to use an example of a custom tool
# from inventory_management.tools.custom_tool import MyCustomTool

# Check our tools documentations for more information on how to use them
# from crewai_tools import SerperDevTool

@CrewBase
class InventoryManagementCrew():
	"""库存管理系统团队"""
	agents_config = 'config/agents.yaml'
	tasks_config = 'config/tasks.yaml'

	def __init__(self, llm: ChatOpenAI = None):
		self.db_tool = DatabaseTool()
		self.nlp_tool = NLPTool(llm)
		self.inventory_tool = InventoryTool(self.db_tool, self.nlp_tool)
		self.llm = llm

	def create_tool(self, name: str, description: str, func) -> Tool:
		"""创建工具"""
		return Tool(
			name=name,
			description=description,
			func=func
		)

	@agent
	def inventory_manager(self) -> Agent:
		return Agent(
			config=self.agents_config['inventory_manager'],
			tools=[
				self.create_tool(
					name="get_low_stock_alerts",
					description="获取库存预警信息",
					func=self.inventory_tool.get_low_stock_alerts
				),
				self.create_tool(
					name="get_inventory_history",
					description="获取库存历史记录",
					func=self.db_tool.get_inventory_history
				),
				self.create_tool(
					name="check_stock",
					description="检查商品库存",
					func=self.db_tool.check_stock
				)
			],
			llm=self.llm,
			verbose=True
		)

	@agent
	def purchase_handler(self) -> Agent:
		return Agent(
			config=self.agents_config['purchase_handler'],
			tools=[
				self.create_tool(
					name="process_restock",
					description="处理补货请求",
					func=self.inventory_tool.process_restock
				),
				self.create_tool(
					name="update_stock",
					description="更新库存",
					func=self.db_tool.update_stock
				),
				self.create_tool(
					name="get_product_by_name",
					description="根据名称查找商品",
					func=self.db_tool.get_product_by_name
				)
			],
			llm=self.llm,
			verbose=True
		)

	@agent
	def sales_handler(self) -> Agent:
		return Agent(
			config=self.agents_config['sales_handler'],
			tools=[
				self.create_tool(
					name="process_purchase",
					description="处理购买请求",
					func=self.inventory_tool.process_purchase
				),
				self.create_tool(
					name="process_payment",
					description="处理支付请求",
					func=self.inventory_tool.process_payment
				),
				self.create_tool(
					name="process_payment_db",
					description="处理支付数据库操作",
					func=self.db_tool.process_payment
				)
			],
			llm=self.llm,
			verbose=True
		)

	@task
	def inventory_task(self) -> Task:
		return Task(
			description="处理库存相关的操作",
			agent=self.inventory_manager,
			expected_output="库存状态报告",
			context=lambda inputs: f"处理库存查询: {inputs.get('query', '')}"
		)

	@task
	def purchase_task(self) -> Task:
		return Task(
			description="处理采购相关的操作",
			agent=self.purchase_handler,
			expected_output="采购处理结果",
			context=lambda inputs: f"处理采购请求: {inputs.get('query', '')}"
		)

	@task
	def sales_task(self) -> Task:
		return Task(
			description="处理销售相关的操作",
			agent=self.sales_handler,
			expected_output="销售处理结果",
			context=lambda inputs: f"处理销售请求: {inputs.get('query', '')}"
		)

	@crew
	def crew(self) -> Crew:
		"""创建库存管理系统团队"""
		return Crew(
			agents=self.agents,
			tasks=self.tasks,
			process=Process.sequential,
			verbose=True
		)